import time
import torch
from torchvision import transforms
from PIL import Image
from src import deeplabv3_resnet50
import json
import numpy as np
import os

'''
用于批量预测语义分割图片
基于deeplabv3_resnet50训练的kitti数据集的脚本
'''


def time_synchronized():
    torch.cuda.synchronize() if torch.cuda.is_available() else None
    return time.time()


def predit_batch():
    aux = False
    classes = 11

    weights_path = "D:\guomengqi\毕业倒计时\课题\语义地图处理部分\deeplab_v3\save_weights\model_565.pth"
    img_folder = "D:\guomengqi\毕业倒计时\课题\语义地图处理部分\deeplab_v3\mapping_img_change\data_change/rgb/"
    # img_list_file = "D:\guomengqi\毕业倒计时\课题\语义地图处理部分\deeplab_v3\mapping_img_change\data_change/namelist_small.txt"
    save_folder = "D:\guomengqi\毕业倒计时\课题\语义地图处理部分\deeplab_v3\mapping_img_change\data_change/"
    palette_path = "D:\guomengqi\毕业倒计时\课题\语义地图处理部分\deeplab_v3\palette.json"

    assert os.path.exists(weights_path), f"未找到权重文件 {weights_path}."
    assert os.path.exists(img_folder), f"未找到图片文件夹 {img_folder}."
    assert os.path.exists(palette_path), f"未找到调色板文件 {palette_path}."

    if not os.path.exists(save_folder + "label_2d_img_raw/"):
        os.makedirs(save_folder + "label_2d_img_raw/")
    if not os.path.exists(save_folder + "label_2d_bin_raw/"):
        os.makedirs(save_folder + "label_2d_bin_raw/")

    with open(palette_path, "rb") as f:
        pallette_dict = json.load(f)
        pallette = []
        for v in pallette_dict.values():
            pallette += v

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))

    # create model
    model = deeplabv3_resnet50(aux=aux, num_classes=classes + 1)

    weights_dict = torch.load(weights_path, map_location='cpu')['model']
    for k in list(weights_dict.keys()):
        if "aux" in k:
            del weights_dict[k]

    model.load_state_dict(weights_dict)
    model.to(device)

    data_transform = transforms.Compose([
        transforms.Resize(520),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
    ])

    for img_name in os.listdir(img_folder):
        img_path = os.path.join(img_folder, img_name)

        saved_color_img_file = save_folder + "label_2d_img_raw/" + img_name.replace(".png", "_color.png")
        saved_bw_img_file = save_folder + "label_2d_img_raw/" + img_name.replace(".png", "_bw.png")
        saved_bin_file = save_folder + "label_2d_bin_raw/" + img_name.replace(".png", ".bin")

        original_img = Image.open(img_path)
        img = data_transform(original_img).unsqueeze(0).to(device)

        model.eval()  # 进入验证模式
        with torch.no_grad():
            img_height, img_width = img.shape[-2:]
            init_img = torch.zeros((1, 3, img_height, img_width), device=device)
            model(init_img)

            t_start = time_synchronized()
            output = model(img.to(device))
            t_end = time_synchronized()
            print("inference time: {}".format(t_end - t_start))

            prediction = output['out'].argmax(1).squeeze(0)
            prediction = prediction.to("cpu").numpy().astype(np.uint8)

            # # 保存彩色图像
            # pred_color = Image.fromarray(prediction)
            # # pred_color.putpalette(pallette)# 转为调色板模式了
            # pred_color = pred_color.convert('RGB')
            # pred_color.save(saved_color_img_file)
            # Load palette from JSON file
            with open(palette_path, "r") as f:
                palette_dict = json.load(f)
            palette = []
            for color in palette_dict.values():
                palette.extend(color)

            color_image = np.zeros((prediction.shape[0], prediction.shape[1], 3), dtype=np.uint8)
            for label, color in palette_dict.items():
                color_image[prediction == int(label)] = color

            color_image = Image.fromarray(color_image)
            color_image.save(saved_color_img_file)

            # Save black and white image
            pred_bw = Image.fromarray(prediction.astype(np.uint8), mode='L')  # Convert to grayscale and save
            pred_bw.save(saved_bw_img_file)

            # Save probability results as binary file
            original_array = np.squeeze(output['out'].softmax(dim=1), axis=0)
            transformed_array = np.transpose(original_array, (1, 2, 0))
            prob_map_flat = transformed_array.squeeze(0).cpu().numpy().astype(
                np.float32).flatten()  # Flatten probability map

            prob_map_flat.tofile(saved_bin_file)  # Save flattened probability map as binary file


if __name__ == '__main__':
    predit_batch()
